# -*- coding: utf-8 -*-
import numpy as np
import matplotlib.pyplot as plt
import scipy.io as sio
from keras.models import Sequential, Model
from keras.layers import Convolution2D, MaxPooling2D, Conv3D, MaxPooling3D, ZeroPadding3D
from keras.layers import Activation, Dropout, Flatten, Dense, BatchNormalization, Input
from keras.utils.np_utils import to_categorical
from sklearn.decomposition import PCA
from keras.optimizers import Adam, SGD, Adadelta, RMSprop, Nadam
import keras.callbacks as kcallbacks
from keras.regularizers import l2
import time
import collections
from sklearn import metrics, preprocessing
import os
from Utils import zeroPadding, normalization, doPCA, modelStatsRecord, averageAccuracy, ssrn_SS_IN

os.environ["CUDA_VISIBLE_DEVICES"] = "0"

def indexToAssignment(index_, Row, Col, pad_length):
    new_assign = {}
    for counter, value in enumerate(index_):
        assign_0 = value // Col + pad_length
        assign_1 = value % Col + pad_length
        new_assign[counter] = [assign_0, assign_1]
    return new_assign

def assignmentToIndex( assign_0, assign_1, Row, Col):
    new_index = assign_0 * Col + assign_1
    return new_index

def selectNeighboringPatch(matrix, pos_row, pos_col, ex_len):
    selected_rows = matrix[range(pos_row-ex_len,pos_row+ex_len+1), :]
    selected_patch = selected_rows[:, range(pos_col-ex_len, pos_col+ex_len+1)]
    return selected_patch

def sampling(proptionVal, groundTruth):              #divide dataset into train and test datasets
    labels_loc = {}
    train = {}
    test = {}
    m = max(groundTruth)
    ksc_sample = [33, 23, 24, 24, 15, 22, 9, 38, 51, 39, 41, 49, 91]
    for i in range(m):
        indices = [j for j, x in enumerate(groundTruth.ravel().tolist()) if x == i + 1]
        np.random.shuffle(indices)
        labels_loc[i] = indices
        nb_val = int(proptionVal * len(indices))
        # train[i] = indices[:-nb_val]
        # test[i] = indices[-nb_val:]
        train[i] = indices[:ksc_sample[i]]
        test[i] = indices[ksc_sample[i]:]
#    whole_indices = []
    train_indices = []
    test_indices = []
    for i in range(m):
#        whole_indices += labels_loc[i]
        train_indices += train[i]
        test_indices += test[i]
    np.random.shuffle(train_indices)
    np.random.shuffle(test_indices)
    return train_indices, test_indices
    
def res4_model_ss():
    model_res4 = ssrn_SS_IN.ResnetBuilder.build_resnet_8((1, img_rows, img_cols, img_channels), nb_classes)

    RMS = RMSprop(lr=0.0003)
    # Let's train the model using RMSprop
    model_res4.compile(loss='categorical_crossentropy', optimizer=RMS, metrics=['accuracy'])

    return model_res4

# mat_data = sio.loadmat('/home/zilong/SSRN/datasets/IN/Indian_pines_corrected.mat')
# data_IN = mat_data['indian_pines_corrected']
# mat_gt = sio.loadmat('/home/zilong/SSRN/datasets/IN/Indian_pines_gt.mat')
#gt_IN = mat_gt['indian_pines_gt']

KSC_dict = sio.loadmat('/data/di.wang/ordinary/23DCNN/KSC.mat')
KSC_gt_dict = sio.loadmat('/data/di.wang/ordinary/23DCNN/KSC_gt.mat')


data_KS = KSC_dict['KSC'].reshape(512,614,176)

gt_KS = KSC_gt_dict['KSC_gt'].reshape(512,614, 1)

print (data_KS.shape)

#new_gt_IN = set_zeros(gt_IN, [1,4,7,9,13,15,16])
new_gt_KS = gt_KS

batch_size = 16
nb_classes = 13
nb_epoch = 200
img_rows, img_cols = 7, 7
patience = 200

INPUT_DIMENSION_CONV = 176
INPUT_DIMENSION = 176

# 20%:10%:70% data for training, validation and testing

TOTAL_SIZE = 5211
VAL_SIZE = 459

TRAIN_SIZE = 459
TEST_SIZE = TOTAL_SIZE - TRAIN_SIZE
VALIDATION_SPLIT = 0.9                     # 20% for trainnig and 80% for validation and testing
# TRAIN_NUM = 10
# TRAIN_SIZE = TRAIN_NUM * nb_classes
# TEST_SIZE = TOTAL_SIZE - TRAIN_SIZE
# VAL_SIZE = TRAIN_SIZE

img_channels = 176
PATCH_LENGTH = 3                #Patch_size (13*2+1)*(13*2+1)

data = data_KS.reshape(np.prod(data_KS.shape[:2]),np.prod(data_KS.shape[2:]))
gt = new_gt_KS.reshape(np.prod(new_gt_KS.shape[:2]),)

#data = preprocessing.scale(data)

scaler = preprocessing.MinMaxScaler(feature_range=(-1,1))
data = scaler.fit_transform(data)

data_ = data.reshape(data_KS.shape[0], data_KS.shape[1],data_KS.shape[2])
whole_data = data_
padded_data = zeroPadding.zeroPadding_3D(whole_data, PATCH_LENGTH)

ITER = 10
CATEGORY = 13

train_data = np.zeros((TRAIN_SIZE, 2*PATCH_LENGTH + 1, 2*PATCH_LENGTH + 1, INPUT_DIMENSION_CONV))
test_data = np.zeros((TEST_SIZE, 2*PATCH_LENGTH + 1, 2*PATCH_LENGTH + 1, INPUT_DIMENSION_CONV))

KAPPA_RES_SS4 = []
OA_RES_SS4 = []
AA_RES_SS4 = []
TRAINING_TIME_RES_SS4 = []
TESTING_TIME_RES_SS4 = []
ELEMENT_ACC_RES_SS4 = np.zeros((ITER, CATEGORY))

seeds = [1220, 1221, 1222, 1223, 1224, 1225, 1226, 1227, 1228, 1229]

#seeds = [1334]

for index_iter in range(ITER):
    print("# %d Iteration" % (index_iter + 1))

    # save the best validated model 
    
    best_weights_RES_path_ss4 = '/data/di.wang/ordinary/23DCNN/SSRN-master/models/KSC_best_RES_3D_SS4_10_' + str(
        index_iter + 11) + '.hdf5'

    np.random.seed(seeds[index_iter])
#    train_indices, test_indices = sampleFixNum.samplingFixedNum(TRAIN_NUM, gt)
    train_indices, test_indices = sampling(VALIDATION_SPLIT, gt)

    # TRAIN_SIZE = len(train_indices)
    # print (TRAIN_SIZE)
    #
    # TEST_SIZE = TOTAL_SIZE - TRAIN_SIZE - VAL_SIZE
    # print (TEST_SIZE)

    y_train = gt[train_indices] - 1
    y_train = to_categorical(np.asarray(y_train))

    y_test = gt[test_indices] - 1
    y_test = to_categorical(np.asarray(y_test))

    # print ("Validation data:")
    # collections.Counter(y_test_raw[-VAL_SIZE:])
    # print ("Testing data:")
    # collections.Counter(y_test_raw[:-VAL_SIZE])

    train_assign = indexToAssignment(train_indices, whole_data.shape[0], whole_data.shape[1], PATCH_LENGTH)
    for i in range(len(train_assign)):
        train_data[i] = selectNeighboringPatch(padded_data, train_assign[i][0], train_assign[i][1], PATCH_LENGTH)

    test_assign = indexToAssignment(test_indices, whole_data.shape[0], whole_data.shape[1], PATCH_LENGTH)
    for i in range(len(test_assign)):
        test_data[i] = selectNeighboringPatch(padded_data, test_assign[i][0], test_assign[i][1], PATCH_LENGTH)

    x_train = train_data.reshape(train_data.shape[0], train_data.shape[1], train_data.shape[2], INPUT_DIMENSION_CONV)
    x_test_all = test_data.reshape(test_data.shape[0], test_data.shape[1], test_data.shape[2], INPUT_DIMENSION_CONV)

    x_val = x_test_all[-VAL_SIZE:]
    y_val = y_test[-VAL_SIZE:]

    x_test = x_test_all[:-VAL_SIZE]
    y_test = y_test[:-VAL_SIZE]

    # SS Residual Network 4 with BN
    model_res4_SS_BN = res4_model_ss()

    earlyStopping6 = kcallbacks.EarlyStopping(monitor='val_loss', patience=patience, verbose=1, mode='auto')
    saveBestModel6 = kcallbacks.ModelCheckpoint(best_weights_RES_path_ss4, monitor='val_loss', verbose=1,
                                                save_best_only=True,
                                                mode='auto')
    tic6 = time.clock()
    print(x_train.shape, x_test.shape)
    history_res4_SS_BN = model_res4_SS_BN.fit(
        x_train.reshape(x_train.shape[0], x_train.shape[1], x_train.shape[2], x_train.shape[3], 1), y_train,
        validation_data=(x_val.reshape(x_val.shape[0], x_val.shape[1], x_val.shape[2], x_val.shape[3], 1), y_val),
        batch_size=batch_size,
        nb_epoch=nb_epoch, shuffle=True, callbacks=[earlyStopping6, saveBestModel6])
    toc6 = time.clock()

    tic7 = time.clock()
    loss_and_metrics_res4_SS_BN = model_res4_SS_BN.evaluate(
        x_test.reshape(x_test.shape[0], x_test.shape[1], x_test.shape[2], x_test.shape[3], 1), y_test,
        batch_size=batch_size)
    toc7 = time.clock()

    print('3D RES_SS4 without BN Training Time: ', toc6 - tic6)
    print('3D RES_SS4 without BN Test time:', toc7 - tic7)

    print('3D RES_SS4 without BN Test score:', loss_and_metrics_res4_SS_BN[0])
    print('3D RES_SS4 without BN Test accuracy:', loss_and_metrics_res4_SS_BN[1])

    print(history_res4_SS_BN.history.keys())

    pred_test_res4 = model_res4_SS_BN.predict(
        x_test.reshape(x_test.shape[0], x_test.shape[1], x_test.shape[2], x_test.shape[3], 1)).argmax(axis=1)
    collections.Counter(pred_test_res4)
    gt_test = gt[test_indices] - 1
    overall_acc_res4 = metrics.accuracy_score(pred_test_res4, gt_test[:-VAL_SIZE])
    confusion_matrix_res4 = metrics.confusion_matrix(pred_test_res4, gt_test[:-VAL_SIZE])
    each_acc_res4, average_acc_res4 = averageAccuracy.AA_andEachClassAccuracy(confusion_matrix_res4)
    kappa = metrics.cohen_kappa_score(pred_test_res4, gt_test[:-VAL_SIZE])
    KAPPA_RES_SS4.append(kappa)
    OA_RES_SS4.append(overall_acc_res4)
    AA_RES_SS4.append(average_acc_res4)
    TRAINING_TIME_RES_SS4.append(toc6 - tic6)
    TESTING_TIME_RES_SS4.append(toc7 - tic7)
    ELEMENT_ACC_RES_SS4[index_iter, :] = each_acc_res4

    print("3D RESNET_SS4 without BN training finished.")
    print("# %d Iteration" % (index_iter + 1))

Experiment_results=np.zeros([13+5,12])

OA=np.array(OA_RES_SS4)
AA=np.array(AA_RES_SS4)
kappa=np.array(KAPPA_RES_SS4)
trn_time=np.array(TRAINING_TIME_RES_SS4)
tes_time=np.array(TESTING_TIME_RES_SS4)

Experiment_results[0,-2]=np.mean(OA)
Experiment_results[0,-1]=np.std(OA)
Experiment_results[1,-2]=np.mean(AA)
Experiment_results[1,-1]=np.std(AA)
Experiment_results[2,-2]=np.mean(kappa)
Experiment_results[2,-1]=np.std(kappa)
Experiment_results[3,-2]=np.mean(trn_time)
Experiment_results[3,-1]=np.std(trn_time)
Experiment_results[4,-2]=np.mean(tes_time)
Experiment_results[4,-1]=np.std(tes_time)
Experiment_results[5:,-2]=np.mean(ELEMENT_ACC_RES_SS4.T,axis=1)
Experiment_results[5:,-1]=np.std(ELEMENT_ACC_RES_SS4.T,axis=1)

Experiment_results[5:,:10]=ELEMENT_ACC_RES_SS4.T

sio.savemat('SSRN_ksc.mat',{'data':Experiment_results})

print('Indian pines dataset training finished!')

# print('+++++++++++++++')
# modelStatsRecord.outputStats(KAPPA_RES_SS4, OA_RES_SS4, AA_RES_SS4, ELEMENT_ACC_RES_SS4,
#                              TRAINING_TIME_RES_SS4, TESTING_TIME_RES_SS4,
#                              history_res4_SS_BN, loss_and_metrics_res4_SS_BN, CATEGORY,
#                              '/data/di.wang/ordinary/23DCNN/SSRN-master/records/IN_train_SS_10_100epoch.txt',
#                              '/data/di.wang/ordinary/23DCNN/SSRN-master/records/IN_train_SS_element_10_100epoch.txt')
#
#
